perf: Avoid copying when materializing output in OrderedPartialAggregateStream - #25312

[P2] Preserve dictionary capacity limits before materializing the whole range

Materializing all completed groups before slicing can exceed a dictionary's key capacity even when every input batch is valid. I reproduced this with GROUP BY (sort_col, nested), where nested is Struct<Dictionary<Int8, Utf8>>, input is ordered only on sort_col, and three valid 64-row batches contain 192 distinct nested dictionary values under the same sort key. With batch_size=32, base returns all expected groups; head panics in RowsGroupColumn::rows_to_array with dictionary re-encode during emit: ArrowError(DictionaryKeyOverflowError, ...).

Both EOF and a subsequent sort-key boundary trigger it. Widening the dictionary keys to Int16 makes the head cases pass. A second reproducer with ordinary integer grouping keys and ARRAY_AGG(Dictionary<Int8, Utf8>) also passes on base and returns DictionaryKeyOverflowError on head during state materialization.

Please retain bounded materialization for key/state encodings that cannot represent the whole completed range in one Arrow array, or emit multiple independently representable batches. Slicing after constructing the combined array is too late to avoid the overflow.

Similar to #25312 (comment) , we should be able to avoid slicing in the long term.

Though I don't fully get the issue for dictionary keys, so I don't know if there is something to fix elsewhere to better address the root cause.

To clarify, this failure happens while constructing the combined output array, before the stream reaches slicing.

Each input batch has its own valid dictionary. The reproducer has three batches with 64 distinct strings each (v000–v063, v064–v127, and v128–v191), and each batch's Int8 keys are only 0–63. There are 192 distinct values across the completed ordered range, but one Int8 dictionary can address only 128 non-null values.

With batch_size=32, base materializes 32 groups at a time and constructs an independent dictionary for each output. Head calls take_completed_state_batch() for all 192 groups, so dictionary construction overflows before entering Outputting. For the nested grouping key, the path is RowsGroupColumn::rows_to_array -> encode_array_if_necessary; its expect turns the overflow into a panic.

I reran both EOF and ordered-boundary cases: they pass on base and fail on head; widening the dictionary keys to Int16 makes them pass. The independent ARRAY_AGG(Dictionary<Int8, Utf8>) case also fails during state materialization, with ordinary integer grouping keys and just one value per group.

A broader fix could normalize the output representation with consistent schema changes, or materialization could produce multiple independently representable batches. Simply removing slicing or changing the panic to a returned error would leave the query failure. The bound needs to account for dictionary child cardinality as well as outer row count; 32 groups works for these reproducers, but is not a general bound for nested aggregate states.

[P2] Handle shared output buffers in downstream merge accounting

These small slices retain the entire completed prefix's buffers. SortPreservingMergeExec accounts for the full backing buffers separately for each buffered batch (BatchBuilder::push_batch), while this stream retains its output reservation until the last slice. This introduces a new query failure under a finite memory pool.

I reproduced it with two ordered input partitions -> Partial Aggregate -> SortPreservingMerge(k1) -> Final Aggregate: each partition contains two ordered k1 ranges of 4,096 distinct (Int32, Int32) groups, using COUNT, 64-row input/output batches, and a 1 MiB pool. Base returns all 8,192 groups with the expected count of 2; head returns no rows and fails with ResourcesExhausted from SortPreservingMergeExec. Enabling disk spilling does not help. Changing only the output batch size to 4,096, avoiding these slices while keeping input batches at 64 rows, makes head pass.

The fallback above does not cover this: reserving the partial output succeeds, but the downstream merge reservation subsequently fails. Please address shared-buffer accounting across batches in the merge, or provide bounded independently releasable output chunks/a finite-memory fallback before introducing this output representation.

The slicing is already implemented in other operator/execution-paths like regular partial/final aggregation, otherwise it will introduce performance penalty since the downstream operator have to handle huge batches.

To solve this issue, I believe now we should fix SMJ memory tracking instead (identify sliced batches and avoid double counting)

With chunked memory management, we can later fully avoid the slicing, and solve this underlying problem completely

• [EPIC] Use blocked / chunked memory management in hash aggregation #24704

Agreed, fixing shared-buffer accounting in the merge is a reasonable way to address this, and I confirmed that other aggregation paths already slice their output. The operator in this reproducer is SortPreservingMergeExec.

I tested an isolated prototype that changes only sorts/builder.rs: use RecordBatchMemoryCounter to count unique backing buffers across the currently live batches plus the incoming batch, and recompute that total after consumed batches are pruned. Upstream aggregate and sort-key cursor reservations were unchanged.

All four previously failing cases pass with that change:

• Integer keys, 1 MiB pool: all 8,192 rows, with disk spilling both disabled and enabled.
• 128-byte string keys, 3.5 MiB pool: all 4,096 rows, with disk spilling both disabled and enabled.

Every aggregate count is 2, no spills occur, and all reservations are released on stream drop. This supports the proposed merge-local fix for these reproducers; we do not need to wait for the full blocked-storage redesign to address this finding. The prototype was only a diagnostic: I have not validated broader behavior or the performance of recomputing the live-buffer set.

The unchanged PR head still fails these cases, so the accounting change and regression coverage should land before resolving this thread. Accounting should follow currently live allocations, since a permanent set of previously seen addresses would be unsafe after buffers are dropped and addresses reused.